"""数值线性代数模块
包含矩阵求逆、线性方程组求解等功能
"""
import numpy as np
from scipy.linalg import solve, lstsq
from scipy.sparse.linalg import cg, gmres, bicgstab
from typing import Tuple, Optional, Union

class NumericalLinearAlgebra:
    @staticmethod
    def solve_linear_system(A: np.ndarray, b: np.ndarray, method: str = 'direct') -> np.ndarray:
        """
        求解线性方程组 Ax = b
        :param A: 系数矩阵
        :param b: 右端向量
        :param method: 求解方法 ('direct', 'least_squares')
        :return: 解向量x
        """
        if method == 'direct':
            return solve(A, b)
        elif method == 'least_squares':
            x, residuals, rank, s = lstsq(A, b)
            return x
        else:
            raise ValueError(f"Unknown method: {method}")
    
    @staticmethod
    def conjugate_gradient(A: np.ndarray, b: np.ndarray, x0: Optional[np.ndarray] = None, 
                          tol: float = 1e-6, maxiter: int = None) -> Tuple[np.ndarray, int]:
        """
        共轭梯度法求解对称正定线性系统
        :param A: 系数矩阵（对称正定）
        :param b: 右端向量
        :param x0: 初始猜测
        :param tol: 容差
        :param maxiter: 最大迭代次数
        :return: (解向量, 迭代次数)
        """
        x, info = cg(A, b, x0=x0, tol=tol, maxiter=maxiter)
        return x, info
    
    @staticmethod
    def gmres_solver(A: np.ndarray, b: np.ndarray, x0: Optional[np.ndarray] = None,
                    tol: float = 1e-6, maxiter: int = None) -> Tuple[np.ndarray, int]:
        """
        GMRES方法求解线性系统
        :param A: 系数矩阵
        :param b: 右端向量
        :param x0: 初始猜测
        :param tol: 容差
        :param maxiter: 最大迭代次数
        :return: (解向量, 收敛信息)
        """
        x, info = gmres(A, b, x0=x0, tol=tol, maxiter=maxiter)
        return x, info
    
    @staticmethod
    def bicgstab_solver(A: np.ndarray, b: np.ndarray, x0: Optional[np.ndarray] = None,
                       tol: float = 1e-6, maxiter: int = None) -> Tuple[np.ndarray, int]:
        """
        BiCGSTAB方法求解线性系统
        :param A: 系数矩阵
        :param b: 右端向量
        :param x0: 初始猜测
        :param tol: 容差
        :param maxiter: 最大迭代次数
        :return: (解向量, 收敛信息)
        """
        x, info = bicgstab(A, b, x0=x0, tol=tol, maxiter=maxiter)
        return x, info
    
    @staticmethod
    def gaussian_elimination(A: np.ndarray, b: np.ndarray, pivoting: str = 'partial') -> np.ndarray:
        """
        高斯消元法
        :param A: 系数矩阵
        :param b: 右端向量
        :param pivoting: 主元策略 ('none', 'partial', 'complete')
        :return: 解向量
        """
        n = A.shape[0]
        Ab = np.column_stack([A.copy(), b.copy()])
        
        # 前向消元
        for i in range(n):
            if pivoting == 'partial':
                # 部分主元
                max_row = i + np.argmax(np.abs(Ab[i:, i]))
                if max_row != i:
                    Ab[[i, max_row]] = Ab[[max_row, i]]
            
            # 消元
            for j in range(i + 1, n):
                if Ab[i, i] != 0:
                    factor = Ab[j, i] / Ab[i, i]
                    Ab[j] -= factor * Ab[i]
        
        # 回代
        x = np.zeros(n)
        for i in range(n - 1, -1, -1):
            x[i] = (Ab[i, -1] - np.dot(Ab[i, i+1:n], x[i+1:n])) / Ab[i, i]
        
        return x
    
    @staticmethod
    def least_squares_solve(A: np.ndarray, b: np.ndarray) -> np.ndarray:
        """
        最小二乘求解
        :param A: 系数矩阵
        :param b: 右端向量
        :return: 最小二乘解
        """
        x, residuals, rank, s = lstsq(A, b)
        return x
    
    @staticmethod
    def error_analysis(A: np.ndarray, x_true: np.ndarray, x_computed: np.ndarray) -> dict:
        """
        误差分析
        :param A: 系数矩阵
        :param x_true: 真实解
        :param x_computed: 计算解
        :return: 误差分析结果
        """
        absolute_error = np.linalg.norm(x_true - x_computed)
        relative_error = absolute_error / np.linalg.norm(x_true)
        condition_number = np.linalg.cond(A)
        
        return {
            'absolute_error': absolute_error,
            'relative_error': relative_error,
            'condition_number': condition_number,
            'error_bound': condition_number * relative_error
        }